#!/usr/local/bin/python2.6
# -*- Mode: Python -*-

import sys
from optparse import OptionParser
import re
import os


cigarPattern   = '([0-9]+[MIDNSHP])'
cigarSearch    = re.compile(cigarPattern)
atomicCigarPattern   = '([0-9]+)([MIDNSHP])'
atomicCigarSearch    =  re.compile(atomicCigarPattern)

class refPos (object):
    """
    struct to store the end position of reference CIGAR operations
    """
    def __init__(self, rPos):
        self.rPos = int(rPos)


class queryPos (object):
    """
    struct to store the start and end positions of query CIGAR operations
    """
    def __init__(self, qsPos, qePos, qLen):
        self.qsPos = int(qsPos)
        self.qePos = int(qePos)
        self.qLen  = int(qLen)


class cigarOp (object):
    """
    sturct to store a discrete CIGAR operations
    """
    def __init__(self, opLength, op):
        self.length = int(opLength)
        self.op     = op

class SAM (object):
    """
    __very__ basic class for SAM input.
    """
    def __init__(self, samList = []):
        if len(samList) > 0:
            self.query    = samList[0]
            self.flag     = int(samList[1])
            self.ref      = samList[2]
            self.pos      = int(samList[3])
            self.mapq     = int(samList[4])
            self.cigar    = samList[5]
            self.matRef   = samList[6]
            self.matePos  = int(samList[7])
            self.iSize    = int(samList[8])
            self.seq      = samList[9]
            self.qual     = samList[10]
            self.tags     = samList[11:]#tags is a list of each tag:vtype:value sets
            self.valid    = 1
        else:
            self.valid = 0
            self.query = 'null'

    def extractCigarOps(self):
        if (self.cigar == "*"):
            self.cigarOps = []
        elif (self.flag & 0x0010):
            cigarOpStrings = cigarSearch.findall(self.cigar)
            cigarOps = []
            for opString in cigarOpStrings:
                cigarOpList = atomicCigarSearch.findall(opString)
                # "struct" for the op and it's length
                cigar = cigarOp(cigarOpList[0][0], cigarOpList[0][1])
                # add to the list of cigarOps
                cigarOps.append(cigar)
                self.cigarOps = cigarOps
            cigarOps.reverse()
            ##do in reverse order because negative strand##
        else:
            cigarOpStrings = cigarSearch.findall(self.cigar)
            cigarOps = []
            for opString in cigarOpStrings:
                cigarOpList = atomicCigarSearch.findall(opString)
                # "struct" for the op and it's length
                cigar = cigarOp(cigarOpList[0][0], cigarOpList[0][1])
                # add to the list of cigarOps
                cigarOps.append(cigar)
                self.cigarOps = cigarOps

    def getCigarOps(self):
        return self.cigarOps
    
    def isReverseStrand(self):
        if (self.flag & 0x0010):
            return True
        return False

    def extractTagValue (self, tagID):
        for tag in self.tags:
            tagParts = tag.split(':', 2);
            if (tagParts[0] == tagID):
                if (tagParts[1] == 'i'):
                    return int(tagParts[2]);
                elif (tagParts[1] == 'H'):
                    return int(tagParts[2],16);
                return tagParts[2];
        return None;

def processSam(opts):
    samFile = opts.samFile;
    if samFile == "stdin":
        s = sys.stdin
    else:
        s = open(samFile, 'r')
    blockList = []
    prevSam = SAM()
    inBlock = 0
    for line in s:
        if line[0] == "@":
            continue
        samList = line.strip().split('\t')
        currSam = SAM(samList)
        if currSam.query != prevSam.query:
            makeBedpe(blockList, opts)
            blockList = [currSam]
            prevSam = currSam
        else:
            blockList.append(currSam)
    # Need to process the last block at end of file.
    makeBedpe(blockList, opts)


def calcOverlap(bed1, bed2):
    o = 1 + min(bed1[6], bed2[6]) - max(bed1[5], bed2[5])
    return max(0, o)

totalLines = 0;
def outputBedpe(bed1, bed2, opts):
    # set1 = set(range(bed1[5], bed1[6] + 1))
    # set2 = set(range(bed2[5], bed2[6] + 1))
    # nonOverlap1 = len(set1) - len(set1 & set2)
    # nonOverlap2 = len(set2) - len(set1 & set2)
    overlap = calcOverlap(bed1, bed2)
    nonOverlap1 = (1 + bed1[6] - bed1[5]) - overlap
    nonOverlap2 = (1 + bed2[6] - bed2[5]) - overlap
    minNonOverlap = min(nonOverlap1, nonOverlap2)  
    if (minNonOverlap >= opts.MNO):
        global totalLines;
        totalLines += 1;
        # sys.stderr.write(bed1[3] + ' ' + str(nonOverlap1) + ' ' + str(nonOverlap2) + ' ' + str(overlap) + ' ' + str(minNonOverlap) + '\n')
        print str(bed1[0]) + "\t" + str(bed1[1]) + "\t" + str(bed1[2]) + "\t" + str(bed2[0]) + "\t" + \
              str(bed2[1]) + "\t" + str(bed2[2]) + "\t" + str(bed1[3]) + "\t" + str(bed1[8]+bed2[8]) + "\t" + str(bed1[4]) + \
              "\t"  + str(bed2[4]) + "\t" + str(bed1[5]) + "\t" + str(bed1[6]) + "\t" + \
              str(bed2[5]) + "\t" + str(bed2[6])  + "\t" + str(minNonOverlap)  + "\t" + str(bed1[9]) + "\t" + \
              "MQ1="+str(bed1[7])+":"+"MQ2="+str(bed2[7])+":"+"AS1="+str(bed1[8])+":"+"AS2="+str(bed2[8])
    

def makeBedpe(blockList, opts):
    if (len(blockList) < 2):
        return
    # First put the "primary" alignments into the pBedBlock.
    # And any remaining "secondary" alignments into the sBedBlock.
    pBedBlock=[]
    sBedBlock=[]
    for i in xrange(len(blockList)):
        blockList[i].extractCigarOps()
        ref = calcRefPosFromCigar(blockList[i].cigarOps, blockList[i].pos)
        query = calcQueryPosFromCigar(blockList[i].cigarOps)
        strand = "+"
        if blockList[i].isReverseStrand() == True:
            strand = "-"
        qual = int(blockList[i].mapq)
        score = blockList[i].extractTagValue('AS');
        # seqname, startRefOff, endRefOff, queryId, strand, startQueryOff, endQueryOff, mapQual, AS, queryLen
        bed = [blockList[i].ref, blockList[i].pos - 1, ref.rPos - 1, blockList[i].query, strand, query.qsPos, query.qePos, qual, score, query.qLen]
        yahaStatus = blockList[i].extractTagValue('YS')
        # For non-YAHA input, assume all alignments are primary.
        if (yahaStatus is None or ((yahaStatus & 0x20) != 0)):
            pBedBlock.append([bed])
        else:
            sBedBlock.append(bed)
    # Sort the primaries by starting query offset.
    pBedBlock.sort(cmp=lambda x,y: cmp(x[0][5],y[0][5]))
    # Now place secondaries into the corresponding primary list.
    for i in xrange(len(sBedBlock)):
        maxOverlap = -1;
        maxJvalue = -1;
        for j in xrange(len(pBedBlock)):
            overlap = calcOverlap(sBedBlock[i], pBedBlock[j][0])
            if (overlap > maxOverlap):
                maxOverlap = overlap
                maxJvalue = j
        if (maxOverlap > 0):
            pBedBlock[maxJvalue].append(sBedBlock[i])
    # Now output all pairs from adjoining similarity sets.
    for i in xrange(len(pBedBlock)-1):
        set1 = pBedBlock[i]
	for j in xrange(i+1, len(pBedBlock)):
            set2 = pBedBlock[j];
            # We need to check for the "span" feature.
            # Let the adjacent ones goes through unmolested.
            p1 = set1[0];
            p2 = set2[0];
            # sys.stderr.write(str(p1[6]) + ' ' + str(p2[5]) + '\n');
            if ((opts.span == 0 and j > i+1) or (opts.span > 0 and p1[6] + opts.span -1 < p2[5])):
                break;
            for ii in xrange(len(set1)):
                for jj in xrange(len(set2)):
                    outputBedpe(set1[ii], set2[jj], opts)


def calcRefPosFromCigar(cigarOps, alignmentStart):
    rPos = alignmentStart
    for cigar in cigarOps:
        if cigar.op == 'M':
            rPos += cigar.length
        elif cigar.op == 'D':
            rPos += cigar.length
        elif cigar.op == 'I':
            continue
        elif cigar.op == 'N':
            raise ValueError('Unexpected Cigar Operation')
    d = refPos(rPos)
    return d

def calcQueryPosFromCigar(cigarOps):
    qsPos = 0
    qePos = 0
    qLen  = 0
    # if first op is a H, need to shift start position 
    # the opPosition counter sees if the for loop is looking at the first index of the cigar object    
    opPosition = 0  
    for cigar in cigarOps:
        if opPosition == 0 and (cigar.op == 'H' or cigar.op == 'S'):
            qsPos += cigar.length
            qePos += cigar.length
            qLen  += cigar.length
        elif opPosition > 0 and (cigar.op == 'H' or cigar.op == 'S'):
            qLen  += cigar.length
        elif cigar.op == 'M' or cigar.op == 'I':
            qePos += cigar.length
            qLen  += cigar.length
        # elif cigar.op == 'D' or cigar.op == 'N':
            # Do nothing.
        opPosition += 1
    d = queryPos(qsPos, qePos, qLen);
    return d


def main():

    usage = """%prog -i <samFile or stdin>>

splitReadSamToBedpe

Author: Michael Lindberg, Aaron Quinlan & Ira Hall	
Description: reports split read mappings in a SAM file to bedpe; 
IMPORTANT NOTE: this replaces previous versions splitReadSamToBedPe and splitReadSamToBedPe_ih (Ira completed final version)
OUTPUT: in addition to standard bedpe format, 
col8=sum of alignment scores; 
col11=queryStart1; 
col12=queryEnd1;
col13=queryStart2; 
col14=queryEnd2; 
col15=minNonOverlap between the two alignments; 
col15=query length;
col16=mapping qualities and alignment scores

Modified by GGF on Jan 12, 2012 to handle secondary alignments.
Modified by GGF on Feb 10, 2012 to add minNonOverlap parameter.
Modified by GGF on March 27, 2012 to add span parameter.
    """

    parser = OptionParser(usage)
    parser.add_option("-i", dest="samFile", help="sam filename sorted by read id (not enforced), or standard input (-i stdin)", metavar="FILE")
    parser.add_option("-m", dest="MNO", help="minimum non-overlap to allow in output (default 25)", metavar="MNO", type="int", default=25);
    parser.add_option("-s", dest="span", help="maximum span to allow between reported pairs (default 0)", metavar="SPAN", type="int", default=0);
    (opts, args) = parser.parse_args()

    if opts.samFile is None:
        parser.print_help()
        print
    else:
        try:
            processSam(opts);
        except IOError as err:
            sys.stderr.write("Output " + str(totalLines) + " total lines.\n");
            sys.stderr.write("IOError " + str(err) + "\n");
            return
            
if __name__ == "__main__":
    sys.exit(main())
    (END)
